#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# Copyright Tencent Technologies Co., Ltd. 2023-2050. All rights reserved.
"""制品(制品包)添加构建签名CMS

必要软件：
JAVA
PYTHON
SIGN_PLUGIN_HOME环境及环境变量
7z

传参的用法:
1.结果目录,输出包路径
python artifacts_sign_cms.py -pp "/usr1/workspace/target/tencent_toolbox" -op "/usr1/shared/tencent_toolbox.zip"
2.包路径,输出包路径
python artifacts_sign_cms.py -pp "/usr1/workspace/target/tencent_toolbox.zip" -op "/usr1/shared/tencent_toolbox.zip"
3.包路径,(输出为原路径)
python artifacts_sign_cms.py -pp "/usr1/workspace/target/tencent_toolbox.zip"
4.包匹配路径,(输出为原路径)
python artifacts_sign_cms.py -pp "/usr1/workspace/target/*.zip"
5.包匹配路径 签名路径(填相对包路径即可) 输出包原路径
python artifacts_sign_cms.py -pp "/usr1/workspace/target/*.zip" -sp "toolbox"

pp是必需,sp和op可选
"""

import os
import argparse
import sys
import subprocess
import hashlib
import json
import stat
import shutil
import shlex
from pathlib import Path
import logging as log
import xml.etree.ElementTree as ET

log.basicConfig(level=log.INFO)
default_workspace = r"d:\workspace" if sys.platform == "win32" else "/usr1/workspace"
WORKSPACE = os.getenv("WORKSPACE", default_workspace)
SIGN_PLUGIN_HOME = os.getenv("SIGN_PLUGIN_HOME", os.getenv("SIGNATURE_HOME"))
s = os.sep


def execute_command(cmd: str, working_dir: str = WORKSPACE) -> None:
    """仅执行系统命令,无返回结果,有报错机制,采用shlex.split所以windows路径请用双引号引起来"""
    log.info("execute current_dir: %s", working_dir)

    if sys.platform == "win32":
        shell_prefix = ["cmd", "/c"]
        cmd_split = shlex.split(cmd)
        command = shell_prefix + cmd_split
    else:
        command = ["bash", "-c", f"{cmd}"]

    log.info("execute command: %s", cmd)

    popen_obj = subprocess.Popen(
        command, cwd=working_dir, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
    )
    while popen_obj.poll() is None:
        line = popen_obj.stdout.readline().strip()
        if line:
            log.info(line.decode("utf8", "ignore"))
    if not popen_obj.returncode:
        log.info("Executed Successfully! About:%s: %s\n", working_dir, cmd)
    else:
        log.info("\n")
        raise Exception(f"Command return code {popen_obj.returncode}")


def parse_args():
    parser = argparse.ArgumentParser(description="This Script is for cms in Package")
    parser.add_argument("--exe_pkg_path", "-pp", required=True, help="Artifact Path After Build")
    parser.add_argument("--sign_relative_path", "-sp", help="Sign Path relative Package Root")
    parser.add_argument("--output_path", "-op", help="Output Path After CMS")
    args = parser.parse_args()
    return args


def obtain_cms_root(directory):
    dir_path = Path(directory)
    num = 0
    for f in dir_path.glob("*"):
        if f.is_dir:
            num += 1
    if num == 1:
        return str(f)
    else:
        return directory


def cms_prepare(arg_dict):
    """做前面前的准备"""

    # 第一个参数：要么是包路径 要么是执行打包的路径(请自行准备：如果是仅打包某一个目录，准备仅包含此目录的文件夹路径)
    exe_pkg_path = Path(arg_dict.exe_pkg_path).absolute()

    exe_pkg_pathlist = []
    sign_root_pathlist = []
    output_pathlist = []

    try:
        if exe_pkg_path.is_dir():
            # 第二个参数：签名根路径（包根路径的相对路径，可不传入，默认为执行打包（包）的根路径）
            sign_root_path = (
                exe_pkg_path
                if not arg_dict.sign_relative_path
                else Path(f"{exe_pkg_path}/{arg_dict.sign_relative_path}")
            )

            # 第三个参数：输出包路径
            output_path = Path(arg_dict.output_path).absolute()

            exe_pkg_pathlist.append(str(exe_pkg_path))
            sign_root_pathlist.append(str(sign_root_path))
            output_pathlist.append(str(output_path))
            return (exe_pkg_pathlist, sign_root_pathlist, output_pathlist)
    except Exception:
        log.info("exe_pkg_path in not dir")

    # 在当前目录中匹配给定的模式
    for pp in exe_pkg_path.parent.glob(str(exe_pkg_path.name)):
        # 如果是包路径，那么创建新的目录用作解压
        file_name = pp.name.rsplit(".", 1)[0]
        exe_pkg_path1 = Path(f"{WORKSPACE}/ws_cms_{file_name}")
        if exe_pkg_path1.exists():
            shutil.rmtree(str(exe_pkg_path1))
        exe_pkg_path1.mkdir(parents=True, exist_ok=True)

        # sign_root_path
        sign_root_path = (
            exe_pkg_path1 if not arg_dict.sign_relative_path else Path(f"{exe_pkg_path1}/{arg_dict.sign_relative_path}")
        )

        # 解压到执行路径
        cmd = ""
        if str(pp).endswith(".zip"):
            # zip解压
            cmd = f'7z x -y "{pp}" -o"{exe_pkg_path1}"'
        if str(pp).endswith(".tar.gz") or str(pp).endswith(".tgz"):
            # gzip和tar解压
            cmd = f'7z x -y "{pp}" -so | 7z x -y -si -ttar -o"{exe_pkg_path1}"'
        execute_command(cmd)

        exe_pkg_pathlist.append(str(exe_pkg_path1))
        sign_root_pathlist.append(str(sign_root_path))
        output_pathlist.append(str(pp))

    return (exe_pkg_pathlist, sign_root_pathlist, output_pathlist)


def sign_sha256file(sha256sync_file):
    """签上CMS"""
    stat_flag = stat.S_IWUSR | stat.S_IRUSR
    common_json = os.path.join(WORKSPACE, "Build", "conf", "COMMON.json")
    with os.fdopen(os.open(common_json, os.O_RDWR, stat_flag), "r") as foo:
        conf_dict = json.loads(foo.read())

    sha256sync_file_dir = Path(sha256sync_file).parent
    sha256sync_name = Path(sha256sync_file).name
    signxml_path = Path(f"{WORKSPACE}/sign_conf.xml")

    # 兼容jenkins签名
    if os.getenv("CLOUD_BUILD2_URL") is None:
        sign_proxy = conf_dict["SIGN_PROXY_JENKINS"]
    else:
        sign_proxy = conf_dict["SIGN_PROXY_CLOUDBUILD"]

    signxml = f"""
<signtasks>
    <signtask name="sign_cms">
    <alias>CMS_Computing_RSA2048_CN_20220810_Tencent</alias>
    <file>{sha256sync_file}</file>
    <crlfile>{sha256sync_file_dir}{s}{sha256sync_name}.crl</crlfile>
    <hashtype>2</hashtype>
    <proxylist>{sign_proxy}</proxylist>
    <signaturestandard>5</signaturestandard>
    <productlineid>049944</productlineid>
    <versionid>{conf_dict["SIGN_VERSIONID"]}</versionid>
    </signtask>
</signtasks>
"""
    signxml_element = ET.fromstring(signxml)
    tree = ET.ElementTree(signxml_element)
    tree.write(str(signxml_path), encoding="utf-8", xml_declaration=True)
    # 执行签名
    execute_command(f'java -jar signature.jar "{signxml_path}"', SIGN_PLUGIN_HOME)


def compress_cmsroot(exe_pkg_p, output_p):
    """压缩到指定位置"""
    cmd = ""
    if Path(output_p).exists():
        Path(output_p).unlink()

    if output_p.endswith(".zip"):
        # zip压缩
        cmd = f'7z a -tzip "{output_p}" *'

    if output_p.endswith(".tar.gz") or output_p.endswith(".tgz"):
        # gzip压缩和tar归档
        p = Path(output_p)
        tar_path = p.stem.split(".", 1)[0]
        cmd = f'7z a -ttar "{tar_path}.tar" * -so | 7z a -tgzip "{output_p}" -si'
    execute_command(cmd, exe_pkg_p)


def encrypt(fpath: str, algorithm: str) -> str:
    with open(fpath, "rb") as f:
        return hashlib.new(algorithm, f.read()).hexdigest()


def sha256sum_sync(cms_root) -> str:
    """遍历每个文件生成sha256sum_sync文件"""
    p = Path(cms_root)
    f_sha256sum = Path(f"{cms_root}/sha256sum_sync")
    with open(str(f_sha256sum), "w+") as f_obj:
        for f_path in p.rglob("*"):
            if str(f_path) == str(f_sha256sum) or ".git" in str(f_path):
                continue
            if f_path.is_file():
                sha256_value = encrypt(str(f_path), "sha256")
                f_rela_path = str(f_path).replace(str(p) + s, "")
                f_obj.write(f"{f_rela_path} {sha256_value}\n")
    return str(f_sha256sum)


def main():
    arg_dict = parse_args()
    exe_pkg_pathlist, sign_root_pathlist, output_pathlist = cms_prepare(arg_dict)
    log.info(f"\n{exe_pkg_pathlist}\n{sign_root_pathlist}\n{output_pathlist}")
    for p, s, o in zip(exe_pkg_pathlist, sign_root_pathlist, output_pathlist):
        sha256sync_path = sha256sum_sync(s)  # sha256_sync文件生成
        sign_sha256file(sha256sync_path)  # 给sha256_sync签名
        compress_cmsroot(p, o)  # 打压缩


if __name__ == "__main__":
    log.info("====================CMS Sign inner_Package Start====================\n")
    main()
    log.info("====================CMS Sign inner_Package End====================\n")
